// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <memory>
#include <vector>

#include <nebula/array/builder_base.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>
#include <turbo/base/checked_cast.h>
#include <turbo/base/macros.h>

namespace nebula {

    class MemoryPool;

    class RecordBatch;

    /// \class RecordBatchBuilder
    /// \brief Helper class for creating record batches iteratively given a known
    /// schema
    class TURBO_EXPORT RecordBatchBuilder {
    public:
        /// \brief create and initialize a RecordBatchBuilder
        /// \param[in] schema The schema for the record batch
        /// \param[in] pool A MemoryPool to use for allocations
        /// \return the created builder instance
        static turbo::Result<std::unique_ptr<RecordBatchBuilder>> create(
                const std::shared_ptr<Schema> &schema, MemoryPool *pool);

        /// \brief create and initialize a RecordBatchBuilder
        /// \param[in] schema The schema for the record batch
        /// \param[in] pool A MemoryPool to use for allocations
        /// \param[in] initial_capacity The initial capacity for the builders
        /// \return the created builder instance
        static turbo::Result<std::unique_ptr<RecordBatchBuilder>> create(
                const std::shared_ptr<Schema> &schema, MemoryPool *pool, int64_t initial_capacity);

        /// \brief Get base pointer to field builder
        /// \param i the field index
        /// \return pointer to ArrayBuilder
        ArrayBuilder *get_field(int i) { return raw_field_builders_[i]; }

        /// \brief Return field builder casted to indicated specific builder type
        /// \param i the field index
        /// \return pointer to template type
        template<typename T>
        T *get_field_as(int i) {
            return turbo::checked_cast<T *>(raw_field_builders_[i]);
        }

        /// \brief finish current batch and optionally reset
        /// \param[in] reset_builders the resulting RecordBatch
        /// \return the resulting RecordBatch
        turbo::Result<std::shared_ptr<RecordBatch>> flush(bool reset_builders);

        /// \brief finish current batch and reset
        /// \return the resulting RecordBatch
        turbo::Result<std::shared_ptr<RecordBatch>> flush();

        /// \brief Set the initial capacity for new builders
        void set_initial_capacity(int64_t capacity);

        /// \brief The initial capacity for builders
        int64_t initial_capacity() const { return initial_capacity_; }

        /// \brief The number of fields in the schema
        int num_fields() const { return schema_->num_fields(); }

        /// \brief The number of fields in the schema
        std::shared_ptr<Schema> schema() const { return schema_; }

    private:
        TURBO_DISALLOW_COPY_AND_ASSIGN(RecordBatchBuilder);

        RecordBatchBuilder(const std::shared_ptr<Schema> &schema, MemoryPool *pool,
                           int64_t initial_capacity);

        turbo::Status CreateBuilders();

        turbo::Status InitBuilders();

        std::shared_ptr<Schema> schema_;
        int64_t initial_capacity_;
        MemoryPool *pool_;

        std::vector<std::unique_ptr<ArrayBuilder>> field_builders_;
        std::vector<ArrayBuilder *> raw_field_builders_;
    };

}  // namespace nebula
